from utils.data_split import DataSplitUtil
from config.config import *

import pandas as pd
from scipy.stats import ttest_ind, chi2_contingency

from utils.utils import md


def compare_groups_with_group_col(df, group_col='ALN status'):
    results = []

    # --------- 自定义分组 ----------
    if 'Age(years)' in df.columns:
        df['age_group'] = pd.cut(df['Age(years)'], bins=[0, 40, float('inf')], labels=['<=40', '>40'])
    if 'Tumour Size(cm)' in df.columns:
        df['tumor_size_group'] = pd.cut(df['Tumour Size(cm)'], bins=[-1, 2, float('inf')], labels=['<=2cm', '>2cm'])
    if 'Ki67' in df.columns:
        df['Ki67_group'] = pd.cut(df['Ki67'], bins=[-1, 0.5, float('inf')], labels=['<=50%', '>50%'])

    group_all = df
    group_ALNM = df[df[group_col] == 1]

    for col in df.columns:
        if col in exclude_columns:
            continue
        if df[col].dtype == 'object' or df[col].nunique() <= 10:  # 分类变量
            group_all_counts = group_all[col].value_counts(normalize=True).sort_index()
            group_ALNM_counts = group_ALNM[col].value_counts(normalize=True).sort_index()
            for level in sorted(df[col].dropna().unique()):
                val0 = f"{group_all[col].value_counts().get(level, 0)} ({group_all_counts.get(level, 0):.1%})"
                val1 = f"{group_ALNM[col].value_counts().get(level, 0)} ({group_ALNM_counts.get(level, 0):.1%})"
                results.append({
                    "Feature": f"{col}={level}",
                    "All": val0,
                    "ALNM": val1,
                })

        else:  # 连续变量
            mean_all, std_all = group_all[col].mean(), group_all[col].std()
            mean_ALNM, std_ALNM = group_ALNM[col].mean(), group_ALNM[col].std()

            results.append({
                "Feature": col,
                "All": f"{mean_all:.2f}({std_all:.2f})",
                "ALNM": f"{mean_ALNM:.2f}({std_ALNM:.2f})",
            })

    return pd.DataFrame(results)


def summarize_clinical_data(df):
    results = []
    total = len(df)

    for col in df.columns:
        if df[col].dtype == 'object' or df[col].nunique() <= 10:
            counts = df[col].value_counts(dropna=False)
            for val, count in counts.items():
                percent = 100 * count / total
                label = f"{col}: {val}"
                results.append((label, f"{count} ({percent:.1f}%)"))
        else:
            mean = df[col].mean()
            std = df[col].std()
            results.append((col, f"{mean:.2f}({std:.2f})"))

    return dict(results)  # 用字典方便合并



def main():
    # 使用示例
    BCNB_clinic_data = pd.read_csv(opj(base_path, 'data/original_data/BCNB/Clinic/patient-clinical-data-process-withoutlabel.csv'))
    # QL_clinic_data = pd.read_csv(opj(base_path, 'data/original_data/QL/Clinic/patient-clinical-data-process.csv'))
    dp = DataSplitUtil(split_random_state=split_random_state_list[0])
    train_raw, test_raw = dp.get_train_test_df(BCNB_clinic_data)
    train = summarize_clinical_data(train_raw)
    test = summarize_clinical_data(test_raw)

    # 合并为总表

    ordered_features = []
    for col in train_raw.columns:
        if col == exclude_columns[0]:
            continue
        if train_raw[col].dtype == 'object' or train_raw[col].nunique() <= 10:
            for val in train_raw[col].dropna().unique():
                ordered_features.append(f"{col}: {val}")
        else:
            ordered_features.append(col)

    summary_df = pd.DataFrame({
        'Feature': ordered_features,
        'Training cohort': [train.get(f, '') for f in ordered_features],
        'Internal validation cohort': [test.get(f, '') for f in ordered_features],
        # 'External validation cohort' : [QL_clinic_data.get(f, '') for f in ordered_features]
    })

    # 临床特征分布
    md(clinic_analysis_result_path)
    summary_df.to_csv(opj(clinic_analysis_result_path, 'each_cohort.csv'), index=False)

    clinic_data = pd.read_csv(clinic_data_path)
    dp = DataSplitUtil(split_random_state=split_random_state_list[0])
    train_raw, test_raw = dp.get_train_test_df(clinic_data)

    # ✳ 分别统计
    entire = compare_groups_with_group_col(pd.concat([train_raw, test_raw]))

    # ✳ 对齐所有特征（以 Entire 为主）
    features = entire['Feature'].tolist()
    merged = pd.DataFrame({'Feature': features})

    # ✳ 合并横向
    merged['Number of cases'] = entire['All']
    merged['Number of ALNM'] = entire['ALNM']

    # 整体，训练集，测试集上各个亚组特征有无转移的p值分析
    merged.to_csv(opj(clinic_analysis_result_path, 'univariate_table.csv'), index=False)

if __name__ == '__main__':
    main()